/*
书闲2025.09.01提供用例，未能在25.04.1上复现问题
*/
#include <hip/hip_runtime.h>
#include <iostream>
#include <cstdlib>

#define THREADS_PER_BLOCK 512

// 原有的reduce核函数
static __global__ void reduce(float* d_in, double* d_out) {
  int myId = threadIdx.x + blockDim.x * blockIdx.x;
  int tid = threadIdx.x;
  for (int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) {
      d_in[myId] += d_in[myId + s];
    }
    __syncthreads();
  }
  if (tid == 0) {
    d_out[blockIdx.x] = d_in[myId];
  }
}

// 添加多个不同名称的核函数用于hipprof验证
static __global__ void kernel_alpha(float* data, int index) {
  int idx = threadIdx.x + blockIdx.x * blockDim.x;
  if (idx == 0) {
    data[index] = 1.0f;
  }
}

static __global__ void kernel_beta(float* data, int index) {
  int idx = threadIdx.x + blockIdx.x * blockDim.x;
  if (idx == 0) {
    data[index] = 2.0f;
  }
}

static __global__ void kernel_gamma(float* data, int index) {
  int idx = threadIdx.x + blockIdx.x * blockDim.x;
  if (idx == 0) {
    data[index] = 3.0f;
  }
}

static __global__ void kernel_delta(float* data, int index) {
  int idx = threadIdx.x + blockIdx.x * blockDim.x;
  if (idx == 0) {
    data[index] = 4.0f;
  }
}

static __global__ void kernel_epsilon(float* data, int index) {
  int idx = threadIdx.x + blockIdx.x * blockDim.x;
  if (idx == 0) {
    data[index] = 5.0f;
  }
}

// HIP错误检查宏
#define HIP_CHECK(error) \
    if (error != hipSuccess) { \
        std::cerr << "HIP error: " << hipGetErrorString(error) << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
        exit(EXIT_FAILURE); \
    }

/**
 * adding nodes to graphs and mapping dependencies.
 */
static void hipGraphsManual(float* inputVec_h, float* inputVec_d,
                            double* outputVec_d, double* result_d, 
                            size_t inputSize, size_t numOfBlocks) {
  hipStream_t streamForGraph;
  hipGraph_t graph;
  hipGraphNode_t memcpyNode, kernelNode;
  double result_h = 0.0;
  HIP_CHECK(hipStreamCreate(&streamForGraph));
  hipKernelNodeParams kernelNodeParams{};

  HIP_CHECK(hipGraphCreate(&graph, 0));
  
  // 内存拷贝节点
  HIP_CHECK(hipGraphAddMemcpyNode1D(&memcpyNode, graph, nullptr, 0, inputVec_d,
                inputVec_h, sizeof(float) * inputSize, hipMemcpyHostToDevice));
  
  // 第一个reduce核函数
  void* kernelArgs[4] = {reinterpret_cast<void*>(&inputVec_d),
                         reinterpret_cast<void*>(&outputVec_d), &inputSize,
                         &numOfBlocks};
  kernelNodeParams.func = reinterpret_cast<void*>(reduce);
  kernelNodeParams.gridDim = dim3(inputSize / THREADS_PER_BLOCK, 1, 1);
  kernelNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);
  kernelNodeParams.sharedMemBytes = 0;
  kernelNodeParams.kernelParams = reinterpret_cast<void**>(kernelArgs);
  kernelNodeParams.extra = nullptr;

  HIP_CHECK(hipGraphAddKernelNode(&kernelNode, graph, &memcpyNode, 1, &kernelNodeParams));

  // 添加多个不同名称的串行核函数
  const int NUM_EXTRA_KERNELS = 5;
  hipGraphNode_t extraKernelNodes[NUM_EXTRA_KERNELS];
  float* test_data_d;
  HIP_CHECK(hipMalloc(&test_data_d, NUM_EXTRA_KERNELS * sizeof(float)));
  
  // 定义不同核函数的配置
  struct KernelConfig {
    void* function;
    const char* name;
    int value;
  };
  
  KernelConfig kernels[] = {
    {reinterpret_cast<void*>(kernel_alpha), "kernel_alpha", 1},
    {reinterpret_cast<void*>(kernel_beta), "kernel_beta", 2},
    {reinterpret_cast<void*>(kernel_gamma), "kernel_gamma", 3},
    {reinterpret_cast<void*>(kernel_delta), "kernel_delta", 4},
    {reinterpret_cast<void*>(kernel_epsilon), "kernel_epsilon", 5}
  };

  // 添加串行的核函数节点
  hipGraphNode_t prevNode = kernelNode;
  for (int i = 0; i < NUM_EXTRA_KERNELS; i++) {
    hipKernelNodeParams extraParams{};
    int index = i;
    void* extraArgs[2] = {&test_data_d, &index};
    
    extraParams.func = kernels[i].function;
    extraParams.gridDim = dim3(1, 1, 1);
    extraParams.blockDim = dim3(1, 1, 1);
    extraParams.sharedMemBytes = 0;
    extraParams.kernelParams = reinterpret_cast<void**>(extraArgs);
    extraParams.extra = nullptr;

    HIP_CHECK(hipGraphAddKernelNode(&extraKernelNodes[i], graph, &prevNode, 1, &extraParams));
    prevNode = extraKernelNodes[i];
    
    std::cout << "Added kernel: " << kernels[i].name << " to graph" << std::endl;
  }

  // 最终的结果拷贝
  HIP_CHECK(hipGraphAddMemcpyNode1D(&memcpyNode, graph, &prevNode, 1, &result_h,
                                    result_d, sizeof(double), hipMemcpyDeviceToHost));

  hipGraphExec_t graphExec;
  HIP_CHECK(hipGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0));
  
  std::cout << "Launching graph with multiple serial kernels..." << std::endl;
  std::cout << "Expected execution order: reduce -> alpha -> beta -> gamma -> delta -> epsilon" << std::endl;
  
  HIP_CHECK(hipGraphLaunch(graphExec, streamForGraph));
  HIP_CHECK(hipStreamSynchronize(streamForGraph));
  
  // 验证结果
  double result_h_cpu = 0.0;
  for (size_t i = 0; i < inputSize; i++) {
    result_h_cpu += inputVec_h[i];
  }
  
  HIP_CHECK(hipGraphExecDestroy(graphExec));
  HIP_CHECK(hipGraphDestroy(graph));
  HIP_CHECK(hipStreamDestroy(streamForGraph));
  HIP_CHECK(hipFree(test_data_d));
  return;
}

int main() {
  constexpr size_t size = 1 << 12;
  constexpr size_t maxBlocks = 512;
  float *inputVec_d = nullptr, *inputVec_h = nullptr;
  double *outputVec_d = nullptr, *result_d = nullptr;

  HIP_CHECK(hipSetDevice(0));
  
  // 分配主机内存
  inputVec_h = reinterpret_cast<float*>(malloc(sizeof(float) * size));
  if (inputVec_h == nullptr) {
    std::cerr << "Failed to allocate host memory!" << std::endl;
    return EXIT_FAILURE;
  }
  
  // 初始化输入数据
  for (size_t i = 0; i < size; i++) {
    inputVec_h[i] = 1.0f;
  }

  // 分配设备内存
  HIP_CHECK(hipMalloc(&inputVec_d, sizeof(float) * size));
  HIP_CHECK(hipMalloc(&outputVec_d, sizeof(double) * maxBlocks));
  HIP_CHECK(hipMalloc(&result_d, sizeof(double)));
  hipGraphsManual(inputVec_h, inputVec_d, outputVec_d, result_d, size, maxBlocks);

  // 清理资源
  HIP_CHECK(hipFree(inputVec_d));
  HIP_CHECK(hipFree(outputVec_d));
  HIP_CHECK(hipFree(result_d));
  free(inputVec_h);

  std::cout << "\nTest completed!" << std::endl;

  return EXIT_SUCCESS;
}